import argparse
import os
from utils.angr_util import AngrUtil
from tqdm import tqdm
import pickle


def run(args):
    path = args.path
    if not os.path.exists(path):
        raise Exception("Dir Not Found!")

    if not os.path.isdir(path):
        raise Exception("Not A Dir!")

    file_list = os.listdir(path)

    ignore_list = []
    f = open(args.ignore_config, "r")
    for line in f.readlines():
        i = int(line.split(":")[0])
        ignore_list.append(i)
    f.close()

    pbar = tqdm(total=len(file_list))
    graph_data = {}

    if args.graph_data != "None":
        f = open(args.graph_data, 'rb')
        graph_data = pickle.load(f)
        f.close()

    for idx, filename in enumerate(file_list):
        pbar.update(1)

        split_names = filename.split("-")
        if len(split_names) < 2:
            continue
        is_auto = split_names[0] == "auto_schedule"
        if is_auto:
            continue

        if idx in ignore_list:
            continue

        if filename in graph_data:
            continue

        util = AngrUtil("{}/{}".format(path, filename))
        util.identify_function()
        graph_data[filename] = {}
        for function_name in util.function_names:
            g = util.get_function_graph(function_name)

            blocks = util.get_func_blocks_by_name(function_name)
            blocks_dict = util.contruct_block_dict(blocks)
            if args.extend == 1:
                g, blocks_dict = util.try_extend_graph(g, blocks_dict)
            graph_data[filename][function_name] = g

        f = open(args.output, "wb")
        pickle.dump(graph_data, f)
        f.close()

    pbar.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser("Generate the trace from CFG")
    parser.add_argument("-p", "--path", type=str, help="The path of .so data")
    parser.add_argument("-e",
                        "--extend",
                        type=int,
                        default=0,
                        help="The trigger of extending the graph")
    parser.add_argument("-o",
                        "--output",
                        type=str,
                        default="graph.pkl",
                        help="The name of the model")
    parser.add_argument("-i",
                        "--ignore_config",
                        type=str,
                        default="ignore_function.txt",
                        help="The config for ignoring files")
    parser.add_argument("-g",
                        "--graph_data",
                        type=str,
                        default="None",
                        help="The history of graph data")

    args = parser.parse_args()
    run(args)
